# %% Import relevant modules
import pandas as pd
import numpy as np
import warnings
import os
from sklearn.impute import KNNImputer
warnings.filterwarnings("ignore")

# %% Loading data
path = r"./EquityData"
files = os.listdir(path)

# file extensions
market_info_extension = "_Prices.csv"
index_info_extension = "_index_information.csv"
industry_info_extension = "_Industry.csv"
fundamental_info_extension = "_fundamentals.csv"
pricing_factors_info_extension = "_total_factors.csv"
valuation_info_extension = "valuation_information.csv"

# filter file names
list_market_info = [elements for elements in files if market_info_extension in elements]
list_index_info = [elements for elements in files if index_info_extension in elements]
list_industry_info = [elements for elements in files if industry_info_extension in elements]
list_fundamental_info = [elements for elements in files if fundamental_info_extension in elements]
list_pricing_factors_info = [elements for elements in files if pricing_factors_info_extension in elements]
list_valuation_info = [elements for elements in files if valuation_info_extension in elements]

# load files
market_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_market_info]
market_info_df = pd.concat(market_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)

index_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_index_info]
index_info_df = pd.concat(index_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)

industry_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_industry_info]
industry_info_df = pd.concat(industry_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)

pricing_factors_info_df = [pd.read_csv(path + os.sep + elements) for elements in list_pricing_factors_info]
pricing_factors_info_df = pd.concat(pricing_factors_info_df).reset_index(
    drop=True).iloc[:, 1:].sort_values(['time', 'code']).reset_index(drop=True)
pricing_factor_names = pricing_factors_info_df.columns[2:-2]

# marge data frames
meta_data = pd.merge(
    market_info_df,
    index_info_df,
    on=["time", "code"],
    how="inner"
)
meta_data = pd.merge(
    meta_data,
    industry_info_df,
    on=["time", "code"],
    how="inner"
)
meta_data = pd.merge(
    meta_data,
    pricing_factors_info_df,
    on=["time", "code"],
    how="inner"
)

meta_data = meta_data.sort_values(['time', 'code']).reset_index(drop=True)
meta_data.to_csv('./meta_data.csv')
#% 缺失值处理
imputer = KNNImputer(n_neighbors=10) # 生成一个KNN模型
meta_data.replace([np.inf, -np.inf], np.nan, inplace=True) # 将无穷值替换为nan
meta_data_imputer = pd.DataFrame(imputer.fit_transform(meta_data[meta_data.columns[2:]]), \
                                 columns=meta_data.columns[2:], index=meta_data.index) # 使用KNN缺失值填充
print(meta_data_imputer.isnull().mean().sum()) # 打印确认无缺失值